#pragma once

#include "geometry.h"
#include "rng.h"

#include <array>

namespace kp {

    enum BxDFType {
        BSDF_REFLECTION = 1 << 0,
        BSDF_TRANSMISSION = 1 << 1,
        BSDF_DIFFUSE = 1 << 2,
        BSDF_GLOSSY = 1 << 3,
        BSDF_SPECULAR = 1 << 4,
        BSDF_ALL = BSDF_DIFFUSE | BSDF_GLOSSY | BSDF_SPECULAR | BSDF_REFLECTION |
        BSDF_TRANSMISSION,
    };

	class BxDF {
    public:
		virtual ~BxDF() {}
        BxDF(BxDFType type) :type(type) {}
        virtual Spectrum f(const Vec3f& wo, const Vec3f& wi) const = 0;
        virtual Spectrum sample_f(const Vec3f& wo, Vec3f* wi,
            const Point2f& sample, Float* pdf, BxDFType* sampledType = nullptr) const;
        virtual Spectrum rho(const Vec3f& wo, int nSamples, const Point2f* samples) const;
        virtual Spectrum rho(int nSamples, const Point2f* samples1, const Point2f* samples2) const;
        virtual Float pdf(const Vec3f& wo, const Vec3f& wi) const;
        
        bool match(BxDFType flag) const { return flag == type; }

        const BxDFType type;
	};

    class LambertianReflection : public BxDF {
    public:
        LambertianReflection(const Spectrum& refl)
            :BxDF(BxDFType(BSDF_REFLECTION | BSDF_DIFFUSE)),R(refl) {}

        virtual Spectrum f(const Vec3f& wo, const Vec3f& wi) const override { return R * InvPi; }
        virtual Spectrum rho(const Vec3f& wo, int nSamples, const Point2f* samples) const override { return R; }
        virtual Spectrum rho(int nSamples, const Point2f* samples1, const Point2f* samples2) const override { return R; }
    private:
        const Spectrum R;
    };

    class BSDF {
    public:
        void add(Ref<BxDF>bxdf) {
            assert(nBxDFs < MaxBxDFs);
            bxdfs[nBxDFs++] = bxdf;
        }

        Spectrum sample_f(const Vec3f& wo, Vec3f* wi,
            const Point2f& sample, Float* pdf, BxDFType* sampledType = nullptr) const {
            Spectrum f;

            return f;
        }

        Spectrum f(const Vec3f& wo, const Vec3f& wi, BxDFType flags) const  {
            Spectrum f;
            
            return f;
        }

        Spectrum rho(const Vec3f& wo, const RNG& rng, BxDFType flags) const  {
            Spectrum r;
            return r;
        }

        Spectrum rho(const RNG& rng, BxDFType flags) const  { 
            Spectrum r;
            return r;
        }

        size_t size() const { return bxdfs.size(); }

    private:
        int nBxDFs = 0;
        static constexpr int MaxBxDFs = 8;
        std::array<Ref<BxDF>, MaxBxDFs> bxdfs;
    };

}
